# 环节2：文档分割 - 把长文字拆成短片段

# 导入需要的库
from langchain.text_splitter import RecursiveCharacterTextSplitter
# 导入上一环节的文档加载函数
from document_loading import load_pdf_document


def split_documents(documents, chunk_size=500, chunk_overlap=100):
    """
    将文档分割成短片段

    参数:
        documents: 从PDF加载的文档列表（Document对象）
        chunk_size: 每个片段的最大字符数
        chunk_overlap: 片段之间的重叠字符数

    返回:
        split_docs: 分割后的文档片段列表
    """
    # 创建文本分割器对象
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,  # 每个片段的字符数
        chunk_overlap=chunk_overlap,  # 片段间重叠的字符数
        length_function=len,  # 用于计算长度的函数
        separators=["\n\n", "\n", " ", ""]  # 分割时优先使用的分隔符
    )

    # 分割文档
    split_docs = text_splitter.split_documents(documents)

    # 打印分割结果信息
    print(f"文档分割完成！原始文档数: {len(documents)}, 分割后片段数: {len(split_docs)}")
    print(f"每个片段大约 {chunk_size} 字符，片段间重叠 {chunk_overlap} 字符")

    return split_docs


# 主程序：测试文档分割功能
if __name__ == "__main__":
    # 1. 先加载PDF文档（复用第一环节的功能）
    pdf_file_path = "平安保险用户手册.pdf"
    document_pages = load_pdf_document(pdf_file_path)

    # 如果加载成功，则进行分割
    if document_pages:
        # 2. 调用分割函数
        split_documents = split_documents(
            documents=document_pages,
            chunk_size=500,  # 可根据需要调整
            chunk_overlap=100  # 可根据需要调整
        )

        # 3. 展示分割结果的前3个片段，验证效果
        print("\n----- 分割结果预览 -----")
        for i, doc in enumerate(split_documents[:3]):  # 只看前3个片段
            print(f"\n片段 {i + 1}:")
            print(f"页码: {doc.metadata.get('page', '未知')}")
            print(f"长度: {len(doc.page_content)} 字符")
            print("内容:", doc.page_content[:100] + "...")  # 只显示前100字符